import numpy as np
import sys
sys.path.append("./")
from PIL import Image
from torchvision import transforms
from model_zoo.inception_local import inception_v3
from model_zoo.resnet_local import resnet101
from model_zoo.densenet_local import densenet121
from model_zoo.vgg_local import vgg11
from model_zoo.conv_tasnet_local import  ConvTasNet
from model_zoo.segmentation_local import deeplabv3_mobilenet_v3_large
import torch
import pandas as pd
import os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
import pickle
from mps_controller import MPSSet
import time
filename = "image/panda.jpg"
layer_name = {"inception_v3":['input', 'Conv2d_1a_3x3', 'Conv2d_2a_3x3', 'Conv2d_2b_3x3', 'Conv2d_3b_1x1', 'Conv2d_4a_3x3',
              'Mixed_5b', 'Mixed_5c', 'Mixed_5d', 'Mixed_6a', 'Mixed_6b', 'Mixed_6c', 'Mixed_6d', 'Mixed_6e',
              'Mixed_7a', 'Mixed_7b', 'Mixed_7c','prediction'],
              "resnet_101": ['input','conv1','layer_1_0', 'layer_1_1', 'layer_1_2', 'layer_2_0', 'layer_2_1', 'layer_2_2',
                            'layer_2_3', 'layer_3_0', 'layer_3_1', 'layer_3_2', 'layer_3_3', 'layer_3_4', 'layer_3_5',
                            'layer_3_6','layer_3_7','prediction'],

              "densenet121":['input','orderdict','block1','trans1','block2','trans2','block3','trans3','prediction'],
              "vgg11":['input','f0','f1','f2','f3','f4','prediction'],
              'mobilenet_v3':['input','conv1','b0', 'b1', 'b2', 'b3', 'b4', 'b5', 'b6','b7', 'b8', 'b9', 'b10', 'b11', 'b12', \
                              'b13', 'b14','conv2','prediction'],
              'convtasnet': ['input', 'group1', 'conv1', 's_0_l_0',  's_0_l_4', 's_1_l_0', 's_1_l_4',  's_2_l_0',
                             's_2_l_4', 'conv2', 'prediction']
              }
def get_model_by_name(model_name, partition_layer,end_layer):
    model = None
    if model_name == "inception_v3":
        model = inception_v3(pretrained=True, progress=True, partition_layer=partition_layer, end_layer=end_layer)
    elif model_name == "resnet_101":
        model = resnet101(partition_layer=partition_layer, end_layer=end_layer)
    elif model_name == "vgg11":
        model = vgg11(partition_layer=partition_layer, end_layer=end_layer)
    elif model_name == 'densenet121':
        model = densenet121(partition_layer=partition_layer, end_layer=end_layer)
    elif model_name == 'mobilenet_v3':
        model = deeplabv3_mobilenet_v3_large(partition_layer=partition_layer, end_layer=end_layer)
    elif model_name == 'convtasnet':
        model = ConvTasNet(partition_layer=partition_layer, end_layer=end_layer)
    return model

def get_slice_time(model_name,device,start_layer):
    if model_name not in layer_name.keys():
        print("model name is invalid", model_name)
        return

    model,input_np,partition_layer_list,end_layer_list = None, None,None, None
    if device in ['nano','edge']:
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        deploy_device = 'cuda:0'

    else:
        deploy_device = 'cpu:0'

    # 1. depend on the type of the device to set the partition layer and the end layer
    if device in ["pi",'nano']:
        if device == 'nano':
            torch.set_num_threads(4)
        samples = 50
        MIN_INPUT_BATCH = 1
        MAX_INPUT_BATCH = 1
        partition_layer_list = ['input']
        end_layer_list = layer_name[model_name][1:]
    elif device == "edge":
        MIN_INPUT_BATCH = 1
        MAX_INPUT_BATCH = 15
        samples = 60
        partition_layer_list = layer_name[model_name][start_layer:-1]
        end_layer_list = ['prediction']

    for partition_layer in partition_layer_list:
        for end_layer in end_layer_list:
            # 2. depend on the type, the partition layer and the end layer to create the model instance
            model = get_model_by_name(model_name,partition_layer,end_layer)
            model.eval()
            model.to(deploy_device)
            with torch.no_grad():
                for batch in range(MIN_INPUT_BATCH, MAX_INPUT_BATCH + 1):
                    model_delay = []
                    input_np = np.load("input_data/" + model_name + "/" + model_name + "_" + partition_layer + ".npy")
                    if model_name == 'convtasnet' and end_layer == 'prediction':
                        input_audio = torch.from_numpy(np.load("input_data/" + model_name + "/" + model_name + "_input.npy").repeat(batch, axis=0)).to(deploy_device)

                    input_batch = torch.from_numpy(input_np.repeat(batch, axis=0)).to(deploy_device)
                    for i in range(samples):
                        if model_name == 'convtasnet' and end_layer == 'prediction':
                            model.upload_input(input_audio)
                        if device in ['edge','nano']:
                            start.record()
                            temp = model(input_batch)
                            end.record()
                            torch.cuda.synchronize()
                            delay = start.elapsed_time(end)
                        else:
                            start = time.time()
                            temp = model(input_batch)
                            end = time.time()
                            delay = end-start
                        model_delay.append(delay)


